#!/usr/bin/env python
import pygame as pg
from random import randint,shuffle,seed
from time import time
import sys
try:
	import psyco
	psyco.full()
except ImportError:
	pass

def clamp(x,lower,upper):
	if x > upper:
		x = upper
	elif x < lower:
		x = lower
	return x

class Polygon:
	points = []
	colour = (0,0,0)
	alpha = 255
	surface = None
	haschanged = True
	def __init__(self):
		pass
	def ReColour(self):
		self.colour = (randint(0,255),randint(0,255),randint(0,255))
		self.alpha = randint(0,255)
		self.haschanged = True
	def ReColourDelta(self,delta):
		newred = clamp(self.colour[0]+randint(-delta,delta),0,255)
		newgreen = clamp(self.colour[1]+randint(-delta,delta),0,255)
		newblue = clamp(self.colour[2]+randint(-delta,delta),0,255)
		self.colour = (newred,newgreen,newblue)
		self.alpha = clamp(self.alpha+randint(-delta,delta),0,255)
		self.haschanged = True
	def Generate(self,points,xmax,ymax):
		self.points = [0]*points
		for i in range(0,points):
			x,y = randint(0,xmax),randint(0,ymax)
			self.points[i] = (x,y)
		self.ReColour()
		self.haschanged = True

	def ReShape(self,xmax,ymax,delta):
		for i in range(0,randint(0,4)):
			choice = randint(0,len(self.points)-1)
			x,y = randint(-delta,delta),randint(-delta,+delta)
			x = clamp(x+self.points[choice][0],0,xmax)
			y = clamp(y+self.points[choice][1],0,ymax)
			self.points[choice] = (x,y)
		self.haschanged = True
	def Draw(self,surface,size):
		if (self.haschanged):
			sf = pg.Surface(size)
			sf.fill((0,0,0))
			pg.draw.polygon(sf, self.colour, self.points)
			sf.set_colorkey((0,0,0))
			sf.set_alpha(self.alpha)
			self.surface = sf
			self.haschanged = False
			surface.blit(sf,(0,0))
		else:
			surface.blit(self.surface,(0,0))
		

	def Clone(self):
		e = Polygon()
		e.points = self.points[0:]
		e.colour = self.colour
		e.alpha = self.alpha
		e.haschanged = True
		return e


class Drawing:
	width = 0
	height = 0
	pointsperpolygon = 6
	polygons = []
	surface = None
	surfacedata = None
	def __init__(self,width,height):
		self.width = width
		self.height = height
	def Mutate(self, numMutations):
		for i in range(0,numMutations):
			e = randint(0,2)
			if e == 0:
				choice = randint(0,len(self.polygons)-1)
				self.polygons[choice].ReColourDelta(5)
			elif e == 1:
				choice = randint(0,len(self.polygons)-1)
				self.polygons[choice].ReShape(self.width,self.height,25)
			elif e == 2:
				c1 = randint(0,len(self.polygons)-1)
				c2 = clamp(c1 + randint(-5,5),0,len(self.polygons)-1)
				self.polygons[c1],self.polygons[c2] = self.polygons[c2],self.polygons[c1]
	def Draw(self,surface):
		surface.fill((255,255,255))
		for polygon in self.polygons:
			polygon.Draw(surface,(self.width,self.height))
	def Generate(self,polygons,sides):
		self.polygons = [0]*polygons
		self.pointsperpolygon = sides
		for i in range(0,polygons):
			e = Polygon()
			#e.ReColour()
			e.Generate(self.pointsperpolygon,self.width,self.height)
			self.polygons[i] = e
	def Clone(self):
		d = Drawing(self.width,self.height)
		length = len(self.polygons)
		d.polygons = [0]*length
		for i in range(0,length):
			d.polygons[i] = self.polygons[i].Clone()
		return d

#global state hack for speed
wrange,hrange = None,None

def fitness_function(surface_1,imagedataarray,upper):
	global wrange
	global hrange
	fitness = 0
	pixelarray = pg.PixelArray(surface_1)
	for i in wrange:
		row,imagerow = pixelarray[i],imagedataarray[i]
		for j in hrange:
			cl,cl2 = pg.Color(row[j]),imagerow[j]
			deltar,deltag,deltab = cl[1]-cl2[0],cl[2]-cl2[1],cl[3]-cl2[2]
			#deltag = cl[2]-cl2[1]
			##deltab = cl[3]-cl2[2]
			fitness += (deltar*deltar)+(deltag*deltag)+(deltab*deltab)
		if (fitness > upper):
			return fitness
				
	return fitness
#psyco.bind(fitness_function)


def main():
	if len(sys.argv) < 2:
		print "in order to use type ./genetic.py \"imagename\" where imagename is the image you want to approximate, followed by any arguments"
		return

	#some 'settings'
	numMutations = 1
	numPolygons = 200
	numSides = 6

	#init the pygame, grab the image and load it
	pg.init()
	imagename = sys.argv[1]
	imagesurf = pg.image.load(imagename)

	imsize = imagesurf.get_size()
	newsurf = pg.Surface(imsize)
	newsurf.fill((0,0,0))
	newsurf.blit(imagesurf,(0,0))
	imagesurf = newsurf
	global hrange
	global wrange
	x,y = imsize
	wrange = range(0,x)
	hrange = range(0,y)
	#generate the target image array
	imagedataarray = [0]*imsize[0]
	for i in range(0,imsize[0]):
		col = [0]*imsize[1]
		for j in range(0,imsize[1]):
			col[j] = imagesurf.get_at((i,j))
		imagedataarray[i] = col

	#create a drawing, initialise display output
	d = Drawing(imsize[0],imsize[1])
	d.Generate(numPolygons,numSides)
	width,height = imsize[0]*3,imsize[1]
	pg.display.set_mode((width,height))
	screen = pg.display.get_surface()
	screen.fill((0,0,0))
	screen.blit(imagesurf,(0,0))
	pg.display.update()

	#move d to parent for the upcoming loop, create a child
	parent = d
	child = d.Clone()

	#init the two drawing surfaces
	childSurface = pg.Surface(imsize)
	parentSurface = pg.Surface(imsize)
	drawpos = imsize[0],0
	drawpos2 = imsize[0]*2,0

	#get the target rectanlges for the child and parent - where they will get drawn
	parentRect = parentSurface.get_rect(topleft=drawpos)
	childRect = childSurface.get_rect(topleft=drawpos2)

	#draw the parent and find out how well it matches (badly at this point)
	parent.Draw(parentSurface)
	p = fitness_function(parentSurface,imagedataarray,100000000000000000000)
	newParent = True
	renders = 0
	tnow = time()

	#for i in range(0,1000):
	while 1:
		#mutate and test the child
		child.Mutate(numMutations)
		child.Draw(childSurface)
		c = fitness_function(childSurface,imagedataarray,p)

		#assert parent != child
		#combine the surfaces for display
		if newParent:
			#only refresh the parent if its changed
			screen.fill((0,0,0),parentRect)
			screen.blit(parentSurface,drawpos)
			pg.display.update()

		if c < p:
			#if the child is a better match
			#store the new 'best' match value
			print "delta: " + str(p-c)
			p = c
			#move the child up to being the new parent
			parent = child.Clone()
			parentSurface.fill((0,0,0))
			parent.Draw(parentSurface)
			print "newlo: " + str(p)
			newParent = True
		else:
			#child is no good - begin afresh
			child,newParent = parent.Clone(),False
		
		renders += 1
		if renders % 100 == 0:
			t = time()
			pg.display.set_caption(str(renders/(t-tnow)))
			tnow = t
			renders = 0
		
		if p == 0:
			raw_input('Done! (press enter to close)')
			break
		#sys.stderr.write(str(i) + "\n")

if __name__ == "__main__":
	profiling = 0
	#seed(3)
	if profiling == 1:
		import cProfile,pstats
		prof = cProfile.Profile()
		prof = prof.runctx("main()", globals(), locals())
		stats = pstats.Stats(prof)
		stats.sort_stats("time")  # Or cumulative
		stats.print_stats(800)  # 80 = how many to print
		# The rest is optional.
		stats.print_callees()
		stats.print_callers()
	else: main()

